#! -*- coding:utf-8 -*-
# 句子对分类任务，2021 搜狐校园 数据集
# val_acc: 0.83, test_acc: 0.65
import os,json
import numpy as np
from bert4keras.backend import keras, set_gelu, K
from bert4keras.tokenizers import Tokenizer
from bert4keras.models import build_transformer_model
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding, DataGenerator
from bert4keras.snippets import open
from keras.layers import Dropout, Dense
import pandas as pd
from tqdm import tqdm


set_gelu('tanh')  # 切换gelu版本

maxlen = 128
batch_size = 16
LR = 1e-5 
DR = 0.2

pre_model ='chinese_wobert_L-12_H-768_A-12'
path_drive ='E:/opt/share/pretrained/'+pre_model
model = path_drive
config_path = os.path.join(path_drive , 'bert_config.json')
checkpoint_path = os.path.join(path_drive ,  'bert_model.ckpt')
dict_path = os.path.join(path_drive,  'vocab.txt') 
model_n = 'best_model.weights'

def prejson(input_path):
    tmp = []
    for line in open(input_path, 'r',encoding='utf-8'):
        tmp.append(json.loads(line))

    data = pd.DataFrame(tmp)
    return data 

train_df = pd.DataFrame() 
valid_df = pd.DataFrame() 
test_df = pd.DataFrame() 
for root,dirs,files in os.walk("sohu2021_open_data_clean"): 
    for dir in dirs: 
        data_root = os.path.join(root,dir)
        train = prejson(os.path.join(data_root,'train.txt')) 
        dev = prejson(os.path.join(data_root,'valid.txt'))
        test = prejson(os.path.join(data_root,'test_with_id.txt'))
        train_df = pd.concat([train,train_df], axis=0, ignore_index=True)
        valid_df = pd.concat([dev,valid_df], axis=0, ignore_index=True)
        test_df = pd.concat([test,test_df], axis=0, ignore_index=True)
          

train_df['labelA'] = train_df['labelA'].fillna(0).astype(int)
train_df['labelB'] = train_df['labelB'].fillna(0).astype(int)

valid_df['labelA'] = valid_df['labelA'].fillna(0).astype(int)
valid_df['labelB'] = valid_df['labelB'].fillna(0).astype(int)

train_df['label'] = train_df['labelA'] + train_df['labelB']
valid_df['label'] = valid_df['labelA'] + valid_df['labelB']

train_df.drop(["labelA", "labelB"], axis=1, inplace=True)
valid_df.drop(["labelA", "labelB"], axis=1, inplace=True)


train_data = train_df[['source', 'target', 'label']].values
valid_data = valid_df[['source', 'target', 'label']].values
 
test_data = test_df[['source', 'target', 'id']].values

# 建立分词器
tokenizer = Tokenizer(dict_path, do_lower_case=True)


class data_generator(DataGenerator):
    """数据生成器
    """
    def __iter__(self, random=False):
        batch_token_ids, batch_segment_ids, batch_labels = [], [], []
        for is_end, (text1, text2, label) in self.sample(random):
            token_ids, segment_ids = tokenizer.encode(
                text1, text2, maxlen=maxlen
            )
            batch_token_ids.append(token_ids)
            batch_segment_ids.append(segment_ids)
            batch_labels.append([label])
            if len(batch_token_ids) == self.batch_size or is_end:
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids = sequence_padding(batch_segment_ids)
                batch_labels = sequence_padding(batch_labels)
                yield [batch_token_ids, batch_segment_ids], batch_labels
                batch_token_ids, batch_segment_ids, batch_labels = [], [], []


# 加载预训练模型
bert = build_transformer_model(
    config_path=config_path,
    checkpoint_path=checkpoint_path,
    with_pool=True,
    return_keras_model=False,
)

output = Dropout(rate=DR)(bert.model.output)
output = Dense(
    units=2, activation='softmax', kernel_initializer=bert.initializer
)(output)

model = keras.models.Model(bert.model.input, output)  
model.summary() 
              
model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer=Adam(LR),  # 用足够小的学习率
    # optimizer=PiecewiseLinearLearningRate(Adam(5e-5), {10000: 1, 30000: 0.1}),
    metrics=['accuracy'],
)


# 转换数据集
train_generator = data_generator(train_data, batch_size)
valid_generator = data_generator(valid_data, batch_size) 


def evaluate(data):
    total, right = 0., 0.
    for x_true, y_true in data:
        y_pred = model.predict(x_true).argmax(axis=1)
        y_true = y_true[:, 0]
        total += len(y_true)
        right += (y_true == y_pred).sum()
    return right / total


class Evaluator(keras.callbacks.Callback):
    """评估与保存
    """
    def __init__(self):
        self.best_val_acc = 0.

    def on_epoch_end(self, epoch, logs=None):
        val_acc = evaluate(valid_generator)
        if val_acc > self.best_val_acc:
            self.best_val_acc = val_acc
            model.save_weights(model_n) 
        print( 
            u'val_acc: %.5f, best_val_acc: %.5f\n' %
            (val_acc, self.best_val_acc)
        )

def predict_to_file(out_file):
    """预测结果到文件
    """ 
    D = []
    for row in tqdm(test_data):
        token_ids, segment_ids = tokenizer.encode(
            row[0], row[1], maxlen=maxlen
        )
        y_pred = model.predict([[token_ids], [segment_ids]])
        mx_scope = y_pred[0][1]
#         print(mx_scope)
        D.append([row[2],int(round(mx_scope))])
        
    df = pd.DataFrame(D,columns=['id','label']) 
    df.to_csv(out_file,index=None)   
    
if __name__ == '__main__':

    evaluator = Evaluator()

    model.fit(
        train_generator.forfit(),
        steps_per_epoch=len(train_generator),
        epochs=5,
        callbacks=[evaluator]
    )

    model.load_weights(model_n)
    print(u'final test acc: %05f\n' % (evaluate(test_generator)))

else:

    model.load_weights(model_n)  
    
    